import os
import cv2
import numpy as np


def load_image_and_mask(image_path, mask_path):
    image = cv2.imread(image_path)
    mask = cv2.imread(mask_path, 0)  # Read mask as grayscale
    return image, mask


def convert_mask_to_yolo_seg_format(mask, class_mapping):
    height, width = mask.shape
    yolo_labels = []
    unique_values = np.unique(mask)
    unique_values = unique_values[unique_values > 0]  # Exclude background (value 0)
    for value in unique_values:
        binary_mask = np.where(mask == value, 1, 0).astype(np.uint8)
        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for contour in contours:
            # Calculate the bounding box
            x, y, w, h = cv2.boundingRect(contour)
            center_x = (x + w / 2) / width
            center_y = (y + h / 2) / height
            norm_width = w / width
            norm_height = h / height
            class_id = class_mapping.get(value, -1)  # -1 if not found in mapping
            if class_id == -1:
                continue
            # Flatten the segmentation polygon coordinates
            segmentation = contour.flatten().astype(float).tolist()
            for i in range(len(segmentation)):
                segmentation[i] /= width if i % 2 == 0 else height
            yolo_labels.append([class_id, center_x, center_y, norm_width, norm_height] + segmentation)
    return yolo_labels


def save_yolo_seg_labels(yolo_labels, label_path):
    bbox_path=label_path.replace("label","bbox")

    with open(label_path, 'w') as f:
        for label in yolo_labels:
            seg_label=[label[0]]+label[5:]
            label_str = " ".join(str(i) for i in seg_label)
            f.write(label_str + "\n")
    with open(bbox_path, 'w') as f2:
        for label in yolo_labels:
            bbox_label=label[:5]
            label_str = " ".join(str(i) for i in bbox_label)
            f2.write(label_str + "\n")


def process_dataset(image_folder, mask_folder, class_mapping, label_folder):
    for root, dirs, files in os.walk(image_folder):
        for file in files:
            if file.lower().endswith(('.jpg', '.png')):
                image_path = os.path.join(root, file)
                mask_path = os.path.join(mask_folder, file)
                label_path = os.path.join(label_folder, file.split('.')[0] + ".txt")
                image, mask = load_image_and_mask(image_path, mask_path)
                if image is None or mask is None:
                    print(f"Failed to read {image_path} or {mask_path}")
                    continue
                yolo_labels = convert_mask_to_yolo_seg_format(mask, class_mapping)
                save_yolo_seg_labels(yolo_labels, label_path)


if __name__ == "__main__":
    # image_folder = "data_zoo/splitsynbboxpcb/image"
    # mask_folder = "data_zoo/splitsynbboxpcb/cpimask"
    # label_folder = "data_zoo/splitsynbboxpcb/cpilabel"


    image_folder = "data_zoo/pcb_truecolor_subcpjrsbg/images"
    mask_folder = "data_zoo/pcb_truecolor_subcpjrsbg/masks"
    # mask_folder = "data_zoo/pcb_truecolor_subcpjrsbg/cpimask"
    label_folder = "data_zoo/pcb_truecolor_subcpjrsbg/labels"
    bbox_folder="data_zoo/pcb_truecolor_subcpjrsbg/bboxs"
    if not os.path.exists(label_folder):  # 使用 os.path.exists() 检查路径是否存在
        os.makedirs(label_folder)  # 若路径不存在，则尝试创建多级目录
    if not os.path.exists(bbox_folder):  # 使用 os.path.exists() 检查路径是否存在
        os.makedirs(bbox_folder)  # 若路径不存在，则尝试创建多级目录
           
    # class_mapping = {
    #     1: 0,
    #     5: 1
    # }
    class_mapping = {x: x for x in range(180)}
    process_dataset(image_folder, mask_folder, class_mapping, label_folder)
# output bbox is centerx,centery,normw,normh